# -*- coding:utf-8 -*-
import numpy as np
import matplotlib.pyplot as mp

#------------------------------------------------------------------------#
#                         METHODE DE RESOLUTION                          #
#------------------------------------------------------------------------#

def step_euler(y, t, h, f):
    "Une itération pour la méthode d'Euler"
    return h*f(y,t) + y

def step_middle(y, t, h, f):
    "Une itération pour la méthode du point milieu"
    tmp = y + (h/2.)*f(y,t)
    return y + h*f(tmp,t+(h/2.))

def step_heun(y, t, h, f):
    "Une itération pour la méthode de Heun"
    p1 = f(y,t)
    y1 = y + h*p1
    p2 = f(y1,t+h)
    return y + 0.5*h*(p1+p2)

def step_kutta(y, t, h, f):
    "Une itération pour la méthode de Runge Kutta 4"
    p1 = f(y,t)
    y1 = y + 0.5*h*p1
    p2 = f(y1,t+0.5*h)
    y2 = y + 0.5*h*p2
    p3 = f(y2,t + 0.5*h)
    y3 = y + h*p3
    p4 = f(y3,t+h)
    return y + (1./6.)*h*(p1+2*p2+2*p3+p4)

#------------------------------------------------------------------------#
#                          CALCUL DE LA SOLUTION                         #
#------------------------------------------------------------------------#

def meth_n_step(y0, t0, n, h, f, meth = step_euler):
    "Calcule la suite de n point pour une méthode donnée"
    res = [y0]
    RES = [np.linalg.norm(y0)]
    t = t0
    for i in range(n):
        res.append(meth(res[i], t, h, f))
        RES.append(np.linalg.norm(res[len(res)-1]))
        t = t + h
    return (res,RES)

def meth_epsilon(y0, t0, tf, eps, f, meth = step_euler):
    "Calcule la suite de point pour une méthode et un epsilon donnés"
    i = 1
    h = np.linalg.norm(abs(tf-t0))
    last = np.linalg.norm(y0)
    res = meth_n_step(y0,t0,1,h,f,meth)
    while(abs((res[1][len(res)-1]-last)) > eps):
        last = res[1][len(res)-1]
        i = i + 1
        h = abs(tf-t0)/i
        res = meth_n_step(y0,t0,i,h,f,meth)
    return res


#------------------------------------------------------------------------#
#                       COMPARAISON DES METHODES                         #
#------------------------------------------------------------------------#

def test_methd1(t0,y0,nmax,f,sol=lambda x:1e30,legend = "exp"):
    "Comparaison des courbes pour les quatre méthodes"
    mp.clf()
    x = np.arange(0.,3.+3./nmax,3./nmax)
    res = meth_n_step(y0, t0, nmax, 3./nmax, f, step_euler)[1]
    g = mp.plot(x, res, linewidth=1.0)
    res = meth_n_step(y0, t0, nmax, 3./nmax, f, step_middle)[1]
    h = mp.plot(x, res, linewidth=1.0)
    res = meth_n_step(y0, t0, nmax, 3./nmax, f, step_heun)[1]
    i = mp.plot(x, res, linewidth=1.0)
    res = meth_n_step(y0, t0, nmax, 3./nmax, f, step_kutta)[1]
    j = mp.plot(x, res, linewidth=1.0)
    if(sol(0) < 1e30):
        x = np.arange(0.,3.,0.001)
        k = mp.plot(x, sol(x), linewidth=1.0)
        mp.legend((g,h,i,j,k),("Euler","Point milieu","Heun","Runge Kutta",legend))
    else:
        mp.legend((g,h,i,j),("Euler","Point milieu","Heun","Runge Kutta"))
    mp.show()
    
def plot_sol(t0,y0,f):
    "Affiche la solution approximée sur l' intervalle [0 10]"
    mp.clf()
    nmax = 200
    x = np.arange(0.,10.+10./nmax,10./nmax)
    res = meth_n_step(y0, t0, nmax, 10./nmax, f, step_kutta)[1]
    g = mp.plot(x, res, linewidth=1.0)
    mp.show()

def test_d2(t0,y0,f,sol,meth = step_euler):
    "Comparaison des courbes pour les quatre méthodes"
    mp.clf()
    x = []
    y = []
    res = meth_n_step(y0, t0, 5, 2*np.pi/5, f, meth)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    g = mp.plot(x,y, linewidth=1.0)
    x = []
    y = []
    res = meth_n_step(y0, t0, 10, 2*np.pi/10, f, meth)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    h = mp.plot(x,y, linewidth=1.0)
    x = []
    y = []
    res = meth_n_step(y0, t0, 20, 2*np.pi/20, f, meth)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    i = mp.plot(x,y, linewidth=1.0)
    x = []
    y = []
    for l in np.arange(0.,2*np.pi,0.001):
        x = x + [sol(l)[0][0]]
        y = y + [sol(l)[1][0]]
    k = mp.plot(x,y, linewidth=1.0)
    mp.legend((g,h,i,k),("n=5","n=10","n=20","ref"))
    mp.show()
    return 0

def test_methd2(t0,y0,f,sol=lambda x:1e30,legend = "exp"):
    "Comparaison des courbes pour les quatre méthodes en dimension 2"
    mp.clf()
    nmax = 10
    x = []
    y = []
    res = meth_n_step(y0, t0, nmax, 2*np.pi/nmax, f, step_euler)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    g = mp.plot(x,y, linewidth=1.0)
    x = []
    y = []
    res = meth_n_step(y0, t0, nmax, 2*np.pi/nmax, f, step_middle)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    h = mp.plot(x,y, linewidth=1.0)
    x = []
    y = []
    res = meth_n_step(y0, t0, nmax, 2*np.pi/nmax, f, step_heun)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    i = mp.plot(x,y, linewidth=1.0)
    x = []
    y = []
    res = meth_n_step(y0, t0, nmax, 2*np.pi/nmax, f, step_kutta)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    j = mp.plot(x,y, linewidth=1.0)
    x = []
    y = []
    for l in np.arange(0.,2*np.pi,0.001):
        x = x + [sol(l)[0][0]]
        y = y + [sol(l)[1][0]]
    k = mp.plot(x,y, linewidth=1.0)
    mp.legend((g,h,i,j,k),("Euler","Point milieu","Heun","Runge Kutta",legend))
    mp.show()

#------------------------------------------------------------------------#
#                           CHAMP DES VECTEURS                           #
#------------------------------------------------------------------------#

def fieldOfVectors(f):
    "Trace le champ de vecteur pour une équation diférentielle donnée"
    x = np.arange(-2., 2., 0.2)
    y = np.arange(-2., 2., 0.2)
    X = []
    Y = []
    fx = []
    fy = []
    for i in range(len(x)):
        for j in range(len(y)):
            a = f(np.array([[x[i]]]),np.array([[y[j]]]))[0]
            fy = fy + [a/(abs(a)*((np.sqrt((1./a)**2+1))))]
            fx = fx + [1./((np.sqrt(a**2+1)))]
            X = X + [x[i]]
            Y = Y + [y[j]]
    mp.clf()
    mp.quiver(X,Y,fx,fy)
    mp.show()

if (__name__ == "__main__"):
    print "------------------------------------------------------------------------"
    print "                   TESTS POUR PLUSIEURS FONCTIONS                       "
    print "------------------------------------------------------------------------"
    print "---------------------------- En dimension 1 ----------------------------"
    
    print "1. Résolution du problème de Cauchy pour f(x)=exp(x) ..."
    t0 = np.array([[0.]])
    y0 = np.array([[1.]])
    f = lambda X,T: X[0][0]
    test_methd1(t0,y0,2,f,np.exp)#n=2

    print "2. Résolution du problème de Cauchy pour f(x)=exp(arctan(x)) ...\n"
    t0 = np.array([[0.]])
    y0 = np.array([[1.]])
    f = lambda X,T: X[0][0]/(1+T[0][0]**2)
    test_methd1(t0,y0,2,f,lambda x: np.exp(np.arctan(x)),"f")
    
    print "---------------------------- En dimension 2 ----------------------------"

    print "1. Résolution du problème de Cauchy pour f(x)=(cos(x),sin(x))..."
    t0 = np.array([[0.]])
    y0 = np.array([[1.],[0.]])
    f = lambda X,T: np.array([[-X[1][0]],
                              [X[0][0]]])
    test_methd2(t0,y0,f,lambda t: np.array([[np.cos(t)],[np.sin(t)]]))
    test_d2(t0,y0,f,lambda t: np.array([[np.cos(t)],[np.sin(t)]]),step_euler)

    print "2. Champ de vecteur pour f'(x)=-2xf(x)^2...\n"
    #Équation défferentielle y'=-2*x*(y^2)
    f = lambda X,Y: np.array([[-2*X[0][0]*Y[0][0]**2]])
    fieldOfVectors(f)
